import openai
import json
import base64
import os
from typing import Optional

class QwenVLAnalyzer:
    def __init__(self, api_key: str, base_url: str = "https://dashscope.aliyuncs.com/compatible-mode/v1"):
        """
        初始化 Qwen-VL-Max API 分析器
        
        参数:
            api_key (str): 您的 DashScope API 密钥
            base_url (str): API 的基础 URL
        """
        self.api_key = api_key
        self.base_url = base_url
        
        # 配置 OpenAI 客户端
        openai.api_key = api_key
        openai.base_url = base_url
    
    def encode_image_to_base64(self, image_path: str) -> str:
        """
        将图像文件编码为 base64 字符串
        
        参数:
            image_path (str): 图像文件路径
            
        返回:
            str: Base64 编码的图像字符串
        """
        try:
            with open(image_path, "rb") as image_file:
                encoded_string = base64.b64encode(image_file.read()).decode('utf-8')
                return encoded_string
        except FileNotFoundError:
            raise FileNotFoundError(f"未找到图像文件: {image_path}")
        except Exception as e:
            raise Exception(f"编码图像时出错: {str(e)}")
    
    def analyze_image(self, image_path: str, prompt: str = "请详细描述这张图片。") -> dict:
        """
        使用 Qwen-VL-Max API 分析图像
        
        参数:
            image_path (str): 图像文件路径
            prompt (str): 分析的文字提示
            
        返回:
            dict: 包含图像描述的 API 响应
        """
        try:
            # 使用 OpenAI SDK 发送请求
            client = openai.OpenAI(
                api_key=self.api_key,
                base_url=self.base_url,
            )            
            response = openai.ChatCompletion.create(
                model="qwen-vl-max",
                messages=[
                    {
                        "role": "user",
                        "content": [
                            {
                                "type": "text",
                                "text": prompt
                            },
                            {
                                "type": "image_url",
                                "image_url": {
                                    "url": f"data:image/jpeg;base64,{self.encode_image_to_base64(image_path)}"
                                }
                            }
                        ]
                    }
                ],
                max_tokens=1000,
                temperature=0.7,
                top_p=0.8
            )
            
            return response.to_dict()
            
        except Exception as e:
            raise Exception(f"API 请求失败: {str(e)}")
    
    def get_image_description(self, image_path: str, prompt: str = "请详细描述这张图片。") -> str:
        """
        获取图像的简单文字描述
        
        参数:
            image_path (str): 图像文件路径
            prompt (str): 分析的文字提示
            
        返回:
            str: 图像描述文字
        """
        try:
            result = self.analyze_image(image_path, prompt)
            
            # 从响应中提取描述
            if 'choices' in result and len(result['choices']) > 0:
                choice = result['choices'][0]
                if 'message' in choice and 'content' in choice['message']:
                    return choice['message']['content']
            
            # 备用方案：如果结构不同，返回完整响应
            return json.dumps(result, indent=2)
            
        except Exception as e:
            return f"分析图像时出错: {str(e)}"

def main():
    """
    演示图像分析的主函数
    """
    # 从环境变量或用户输入获取 API 密钥
    api_key = os.getenv("DASHSCOPE_API_KEY")
    
    if not api_key:
        print("请设置您的 DASHSCOPE_API_KEY 环境变量")
        print("或手动输入您的 API 密钥:")
        api_key = input("API 密钥: ").strip()
        
        if not api_key:
            print("未提供 API 密钥。退出。")
            return
    
    # 初始化分析器
    analyzer = QwenVLAnalyzer(api_key)
    
    # 使用 img 文件夹中的图像进行示例
    image_path = "img/train1.jpeg"  # 您可以将其更改为 img 文件夹中的任何图像
    
    if not os.path.exists(image_path):
        print(f"未找到图像文件: {image_path}")
        print("img 文件夹中可用的图像:")
        img_folder = "img"
        if os.path.exists(img_folder):
            for file in os.listdir(img_folder):
                if file.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.bmp')):
                    print(f"  - {file}")
        return
    
    print(f"正在分析图像: {image_path}")
    print("-" * 50)
    
    try:
        # 使用自定义提示分析图像
        prompt = "请提供这张图片的详细描述，包括任何物体、人物、动作、颜色和整体场景。"
        
        # 获取描述
        description = analyzer.get_image_description(image_path, prompt)
        
        print("图像描述:")
        print(description)
        
        # 同时获取完整的 API 响应用于调试
        print("\n" + "="*50)
        print("完整 API 响应:")
        full_response = analyzer.analyze_image(image_path, prompt)
        print(json.dumps(full_response, indent=2))
        
    except Exception as e:
        print(f"错误: {str(e)}")

if __name__ == "__main__":
    main()
